from langchain_openai import ChatOpenAI
from langchain_core.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from role_play.domain.charMeta import CharMeta
from langchain.globals import set_debug


class RoleDesigner:
    def __init__(self, model_name: str = 'GLM-4-Plus', verbose: bool = False):
        set_debug(verbose)
        model = ChatOpenAI(model_name=model_name,
                           api_key="01020b6e3f0b99db1641a08661198d4a.8zeaLRZ03xyf949h",
                           base_url="https://open.bigmodel.cn/api/paas/v4/")
        system_prompt = ("你是一名角色设计师，请根据以下文本生成角色的人设，包括角色的姓名、年龄、性别、外貌特征、性格特征、背景故事和梦想等信息。\n"
                         "你回答的格式必须满足以下json格式,你必须确保你的回答可以在我不做任何处理的前提下可以被python的json.loads()解析，"
                         "其中role_info的value是纯字符串，不是json格式:\n"
                         """
                         {{
                            "name": "角色姓名",
                            "role_info": "角色人设信息"
                         }}
                         """
                         )
        human_prompt = "{text}"

        chat_prompt = ChatPromptTemplate.from_messages(
            [
                SystemMessagePromptTemplate.from_template(system_prompt),
                HumanMessagePromptTemplate.from_template(human_prompt)
            ]
        )

        self.chain = chat_prompt | model | JsonOutputParser()

    def invoke(self,user_desc: str, bot_desc: str) -> CharMeta:
        user=self.chain.invoke({'text': user_desc})
        bot =self.chain.invoke({'text': bot_desc})
        return {'user_name': user['name'], 'user_info': user['role_info'], 'bot_name': bot['name'], 'bot_info': bot['role_info']}
